import torch
from torch.nn import Module
from torchaudio.transforms import AmplitudeToDB
from data_loader.utils.io import TorchScaler

class Clamp(Module):
    def __init__(self, min=0, max=1):
        super().__init__()
        self.min = min
        self.max = max

    def forward(self, x):
        return x.clamp(self.min, self.max)

class Amp2dB(Module):
    def __init__(self):
        super().__init__()
        self.amptodb = AmplitudeToDB(stype='amplitude')
        self.amptodb.amin = 1e-5

    def forward(self, x):
        return self.amptodb(x)

class SEDNorm(Module):
    def __init__(self, norm_dim=[1,2]):
        super(SEDNorm, self).__init__()
        self.sed_norm = torch.nn.Sequential(
            Amp2dB(),
            Clamp(min=-50, max=80),
            TorchScaler("instance", "minmax", norm_dim)
        )
    def forward(self, x):
        return self.sed_norm(x)